/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package de.up.ling.irtg.automata;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.SetMultimap;
import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayFIFOQueue;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntIterable;
import it.unimi.dsi.fastutil.ints.IntList;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntPriorityQueue;
import it.unimi.dsi.fastutil.ints.IntSet;
import it.unimi.dsi.fastutil.objects.Object2IntFunction;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * TODO add weights
 *
 * @author koller
 */
public class Determinizer<State> {
    private final static boolean DEBUG = false;
    private Set<Rule> newRules;
    private List<IntSet> newStateToOld;
    private Object2IntFunction<IntSet> oldStatesToNew;
    private int nextNewState;
    private IntPriorityQueue agenda;
    private StateToRuleIndex oldStateToRulesIndex;
    private Int2ObjectMap<IntSet> oldStateToContainingNewStates;    // o(q) = { Q | q in Q }
    private IntSet newFinalStates;
    private TreeAutomaton<State> auto;

    public Determinizer() {
        newRules = new HashSet<Rule>();

        newStateToOld = new ArrayList<IntSet>();
        newStateToOld.add(null); // add newStateToOld(0) so nextNewState is index of first un-added element in this list

        oldStatesToNew = new Object2IntOpenHashMap<IntSet>();
        oldStateToRulesIndex = new StateToRuleIndex();
        oldStateToContainingNewStates = new Int2ObjectOpenHashMap<IntSet>();
        newFinalStates = new IntOpenHashSet();
        nextNewState = 1;
    }

    public TreeAutomaton<List<State>> determinize(TreeAutomaton<State> auto) {
        this.auto = auto;
        agenda = new IntArrayFIFOQueue();

        if(DEBUG) System.err.println(auto.stateInterner);

        // index rules of old automaton so they can be found easily
        for (Rule rule : auto.getRuleSet()) {
            for (int pos = 0; pos < rule.getArity(); pos++) {
                if(DEBUG) System.err.println("indexing " + rule.toString(auto) + " @" + pos + " = " + rule.getChildren()[pos]);
                oldStateToRulesIndex.put(rule.getChildren()[pos], rule, pos);
            }
        }

        if(DEBUG) System.err.println(oldStateToRulesIndex);

        // process nullary symbols and add states for them to the agenda
        for (int sym = 0; sym <= auto.getSignature().getMaxSymbolId(); sym++) {
            if (auto.getSignature().getArity(sym) == 0) {
                IntSet parents = new IntOpenHashSet();

                for (Rule rule : auto.getRulesBottomUp(sym, new int[0])) {
                    parents.add(rule.getParent());
                }

                if (!parents.isEmpty()) {
                    int newParent = addOrFindState(parents);
                    newRules.add(makeRule(newParent, sym, new IntArrayList()));
                }
            }
        }

        // iterate over agenda until it is empty
        TransitionSet transitionsForEachLabel = new TransitionSet();

        while (!agenda.isEmpty()) {
            int newState = agenda.dequeueInt();

            if(DEBUG) System.err.println("\n\npop newstate: " + newState + " = " + newStateToOld.get(newState) + " = " + makeStateSet(newStateToOld.get(newState)));

            transitionsForEachLabel.clear();

            for (int oldState : newStateToOld.get(newState)) {
                if(DEBUG) System.err.println("oldstate: " + oldState + " = " + auto.getStateForId(oldState));

                for (int label : oldStateToRulesIndex.getLabelsForState(oldState)) {
                    if(DEBUG) System.err.println("label: " + label + " = " + auto.getSignature().resolveSymbolId(label));

                    for (Map.Entry<Rule, IntList> entry : oldStateToRulesIndex.getRulesForState(oldState, label).entrySet()) {
                        Rule rule = entry.getKey();
                        for (int pos : entry.getValue()) {
                            collectNewRules(newState, rule, pos, transitionsForEachLabel);
                        }
                    }
                }
            }


            for (int label : transitionsForEachLabel.getLabels()) {
                for (Map.Entry<IntList, IntSet> entry : transitionsForEachLabel.getTransitionsForLabel(label).entrySet()) {
                    int newParent = addOrFindState(entry.getValue()); // adds newParent to agenda if it was unknown
                    Rule newRule = makeRule(newParent, label, entry.getKey());
                    if(DEBUG) System.err.println("added rule: " + newRule);
                    newRules.add(newRule);
                }
            }
        }

        // create output automaton
        TreeAutomaton<List<State>> ret = new ConcreteTreeAutomaton<List<State>>();
        ret.signature = auto.signature;

        for (int i = 1; i < nextNewState; i++) {
            List<State> ss = makeStateSet(newStateToOld.get(i));
            ret.addState(ss);

            if (newFinalStates.contains(i)) {
                ret.addFinalState(i);
            }
        }

        for (Rule rule : newRules) {
            ret.storeRule(rule);
        }

        return ret;
    }

    private static class TransitionSet {

        Int2ObjectMap<SetMultimap<IntList, Integer>> transitionsForEachLabel = new Int2ObjectOpenHashMap<SetMultimap<IntList, Integer>>();

        public void clear() {
            transitionsForEachLabel.clear();
        }

        private SetMultimap<IntList, Integer> getMapForLabel(int label) {
            SetMultimap<IntList, Integer> ret = transitionsForEachLabel.get(label);

            if (ret == null) {
                ret = HashMultimap.create();
                transitionsForEachLabel.put(label, ret);
            }

            return ret;
        }

        public void add(int label, IntList newChildren, int oldParent) {
            getMapForLabel(label).put(newChildren, oldParent);
        }

        public IntSet getLabels() {
            return transitionsForEachLabel.keySet();
        }

        public Map<IntList, IntSet> getTransitionsForLabel(int label) {
            Map<IntList, IntSet> ret = new HashMap<IntList, IntSet>();

            for (Map.Entry<IntList, Collection<Integer>> entry : getMapForLabel(label).asMap().entrySet()) {
                ret.put(entry.getKey(), new IntOpenHashSet(entry.getValue()));
            }

            return ret;
        }
    }

    private List<State> makeStateSet(IntSet states) {
        List<State> originalOldStates = new ArrayList<State>();

        // convert set to list with predictable order to ensure that
        // all copies of state have same string representation
        IntList sortedStates = new IntArrayList(states);
        Collections.sort(sortedStates);

        for (int o : sortedStates) {
            originalOldStates.add(auto.getStateForId(o));
        }

        return originalOldStates;
    }

    private Rule makeRule(int newParent, int label, IntList newChildren) {
        return new Rule(newParent, label, newChildren.toIntArray(), 1); // TODO weight
    }

    private void collectNewRules(int newState, Rule rule, int pos, TransitionSet transitionsForLabel) {
        assert rule.getArity() > 0;

        if (rule.getArity() == 1) {
            assert pos == 0;

            IntList key = new IntArrayList(new int[]{pos});
            transitionsForLabel.add(rule.getLabel(), key, rule.getParent());
        } else {
            assert rule.getArity() == 2; // ASSPT for now: binary
            assert (pos == 0 || pos == 1);
            int otherPos = 1 - pos;

            IntSet containingNewStates = oldStateToContainingNewStates.get(rule.getChildren()[otherPos]);

            // if new states that contain the old state have been discovered already, generate new rules
            if (containingNewStates != null) {
                for (int q : containingNewStates) {
                    int[] children = new int[2];
                    children[pos] = newState;
                    children[otherPos] = q;
                    transitionsForLabel.add(rule.getLabel(), new IntArrayList(children), rule.getParent());
                }
            }
        }
    }

    private static class StateToRuleIndex {

        private Int2ObjectMap<Int2ObjectMap<Map<Rule, IntList>>> index;  // index(q)(f)(rule) = list(pos of q in rule)

        public StateToRuleIndex() {
            index = new Int2ObjectOpenHashMap<Int2ObjectMap<Map<Rule, IntList>>>();
        }

        public IntIterable getLabelsForState(int oldState) {
            Int2ObjectMap<Map<Rule, IntList>> map = index.get(oldState);

            if (map == null) {
                return new IntArrayList();
            } else {
                return index.get(oldState).keySet();
            }
        }

        public Map<Rule, IntList> getRulesForState(int oldState, int label) {
            assert index.containsKey(oldState);
            assert index.get(oldState).containsKey(label);

            return index.get(oldState).get(label);
        }

        public void put(int state, Rule rule, int position) {
            Int2ObjectMap<Map<Rule, IntList>> map = index.get(state);

            if (map == null) {
                map = new Int2ObjectOpenHashMap<Map<Rule, IntList>>();
                index.put(state, map);
            }

            Map<Rule, IntList> mapForLabel = map.get(rule.getLabel());

            if (mapForLabel == null) {
                mapForLabel = new IdentityHashMap<Rule, IntList>();
                map.put(rule.getLabel(), mapForLabel);
            }

            IntList positions = mapForLabel.get(rule);
            if (positions == null) {
                positions = new IntArrayList();
                mapForLabel.put(rule, positions);
            }

            positions.add(position);
        }

        @Override
        public String toString() {
            return index.toString();
        }
    }

    // If oldStates not known so far, make a new state ID for it
    // and add it to the agenda.
    private int addOrFindState(IntSet oldStates) {
        int ret = oldStatesToNew.getInt(oldStates);

        if (ret == 0) {
            // unknown state set; introduce state to new automaton
            ret = nextNewState++;

            if(DEBUG) System.err.println("new state " + ret + ": " + oldStates + " = " + makeStateSet(oldStates));

            oldStatesToNew.put(oldStates, ret);
            newStateToOld.add(oldStates);

            for (int oldState : oldStates) {
                IntSet s = oldStateToContainingNewStates.get(oldState);

                if (s == null) {
                    s = new IntOpenHashSet();
                }

                oldStateToContainingNewStates.put(oldState, s);
                s.add(ret);

                if(DEBUG) System.err.println("os2cns is now " + oldStateToContainingNewStates);

                if (auto.getFinalStates().contains(oldState)) {
                    newFinalStates.add(ret);
                }
            }

            agenda.enqueue(ret);
        }

        return ret;
    }
}
